import pandas as pd
import os
import logging

def save_report(df, file_path, column_mapping=None):
    """
    根据文件扩展名，将DataFrame保存为格式化的Excel或CSV文件。
    
    :param df: 要保存的DataFrame
    :param file_path: 文件路径（支持 .xlsx 和 .csv）
    :param column_mapping: 一个字典，用于将原始列名映射到新的中文列名
    """
    # 检查文件扩展名
    _, ext = os.path.splitext(file_path)

    # 如果提供了列名映射，则创建一个副本并重命名
    if column_mapping:
        df_to_save = df.copy()
        df_to_save.rename(columns=column_mapping, inplace=True)
    else:
        df_to_save = df

    if ext.lower() == '.csv':
        # 保存为CSV文件
        try:
            df_to_save.to_csv(file_path, index=False, encoding='utf-8-sig')
            print(f"报告已成功保存为CSV: {file_path}")
        except Exception as e:
            print(f"保存CSV文件 '{file_path}' 失败: {e}")

    elif ext.lower() == '.xlsx':
        # 使用现有的Excel保存逻辑
        try:
            _save_df_to_excel_formatted_internal(df_to_save, file_path)
            print(f"报告已成功保存为Excel: {file_path}")
        except Exception as e:
            print(f"保存格式化Excel文件 '{file_path}' 失败: {e}")
            # 如果格式化失败，则回退到普通保存
            df_to_save.to_excel(file_path, index=False)
    else:
        print(f"错误：不支持的文件格式 '{ext}'。请使用 .csv 或 .xlsx。")


def _save_df_to_excel_formatted_internal(df_to_save, file_path):
    """
    内部函数：使用 xlsxwriter 引擎保存DataFrame到Excel，并自动调整列宽、居中对齐。
    """
    # 创建一个Excel writer对象
    with pd.ExcelWriter(file_path, engine='xlsxwriter') as writer:
        df_to_save.to_excel(writer, index=False, sheet_name='Sheet1')

        # 获取workbook和worksheet对象
        workbook = writer.book
        worksheet = writer.sheets['Sheet1']

        # 创建一个居中对齐的格式
        center_format = workbook.add_format({'align': 'center', 'valign': 'vcenter'})

        # 定义一个函数来计算字符串的显示宽度（英文1，中文2）
        def get_str_width(s):
            width = 0
            for char in str(s):
                if '\u4e00' <= char <= '\u9fff':
                    width += 2
                else:
                    width += 1
            return width

        # 迭代每一列，设置宽度和格式
        for i, col in enumerate(df_to_save.columns):
            # 计算列头和内容的最大宽度
            column_len = max(
                df_to_save[col].astype(str).map(get_str_width).max(),
                get_str_width(col)
            ) + 2  # 加一点缓冲
            # 设置列宽和居中格式
            worksheet.set_column(i, i, column_len, center_format)

def ping():
    """
    A simple function to be called by the webapp to check communication.
    """
    return {"status": "ok", "message": "Hello from btstrategy-engine!"} 